#include <errno.h>
#include <pthread.h>
#include <string.h>
#include <unistd.h>
#include <sys/poll.h>

#define __USE_GNU
#include <fcntl.h>
#undef __USE_GNU

#include "compat.h"

#include "splice.h"

#include "psd_logger.h"
#include "psd_measurement.h"
#include "psd_session.h"
#include "psd_xfer.h"
#include "psd_modules.h"
#include "psd_backend_handler.h"

#include "psd_conn_tcp.h"

#define SPLICE_SIZE (64*1024)


void *psd_xfer_splice_src(void *arg);
void *psd_xfer_splice_sink(void *arg);
int psd_xfer_splice_init();
static int psd_xfer_splice_main(psdXferArgs *p_args, psdXferArgs *c_args, uint64_t xfer_amt);

psdModule psd_xfer_splice_module = {
        .desc = "Splice Module",
        .dependencies = "tcp",
        .init = psd_xfer_splice_init
};

psdBackendHandler psd_xfer_splice_backend_handler = {
        .name = "splice",
        .run = psd_xfer_splice_main
};

psdModule *module_info() {
        return &psd_xfer_splice_module;
}

int psd_xfer_splice_init() {
        return psd_add_backend_handler(&psd_xfer_splice_backend_handler);
}

static int psd_xfer_splice_main(psdXferArgs *p_args, psdXferArgs *c_args, uint64_t xfer_amt) {
	int ret;
        pthread_t thread1, thread2;
	pthread_t sink_thr1, sink_thr2;

	if ((strcmp("TCP", p_args->src->protocol) != 0) ||
	    (strcmp("TCP", c_args->src->protocol) != 0)) {
		psd_err(0, "Splice backend can only operate on TCP sockets");
		goto error_exit;
	}

	// setup splice pipes
	ret = pipe(c_args->pipes);

	if (ret < 0) {
		psd_err(0, "pipe failed: %s", strerror(errno));
		goto error_exit;
	}


	ret = pipe(p_args->pipes);

	if (ret < 0) {
		psd_err(0, "pipe failed: %s", strerror(errno));
		goto error_exit;
	}
	
	
	// start the sink threads
	pthread_mutex_lock(&(p_args->sink_lock));
	// create an I/O thread
	if (pthread_create(&sink_thr1, NULL, psd_xfer_splice_sink, (void *) p_args) < 0) {
		pthread_mutex_unlock(&(p_args->sink_lock));
		goto error_exit;
	}

	// wait for the child to initialize
	pthread_cond_wait(&(p_args->sink_cond), &(p_args->sink_lock));
	pthread_mutex_unlock(&(p_args->sink_lock));


	pthread_mutex_lock(&(c_args->sink_lock));
	// create an I/O thread
	if (pthread_create(&sink_thr2, NULL, psd_xfer_splice_sink, (void *) c_args) < 0) {
		pthread_mutex_unlock(&(c_args->sink_lock));
		goto error_exit1;
	}

	// wait for the child to initialize
	pthread_cond_wait(&(c_args->sink_cond), &(c_args->sink_lock));
	pthread_mutex_unlock(&(c_args->sink_lock));

	// create "src" threads
	
        // create an I/O thread
        pthread_mutex_lock(&(p_args->lock));
        if (pthread_create(&thread1, NULL, psd_xfer_splice_src, (void *) p_args) < 0) {
                pthread_mutex_unlock(&(p_args->lock));
                goto error_exit2;
        }

        // wait for the child to initialize
        pthread_cond_wait(&(p_args->cond), &(p_args->lock));
        pthread_mutex_unlock(&(p_args->lock));

        // create an I/O thread
        pthread_mutex_lock(&(c_args->lock));
        if (pthread_create(&thread2, NULL, psd_xfer_splice_src, (void *) c_args) < 0) {
                pthread_mutex_unlock(&(c_args->lock));
                goto error_exit3;
        }
        // wait for the child to initialize
        pthread_cond_wait(&(c_args->cond), &(c_args->lock));
        pthread_mutex_unlock(&(c_args->lock));

        // setup the threads to transfer the specified amount
        pthread_mutex_lock(&(p_args->lock));
        pthread_mutex_lock(&(c_args->lock));
        {
                p_args->xfer_amt = xfer_amt;
                pthread_cond_signal(&(p_args->cond));

                c_args->xfer_amt = xfer_amt;
                pthread_cond_signal(&(c_args->cond));
        }
        pthread_mutex_unlock(&(p_args->lock));
        pthread_mutex_unlock(&(c_args->lock));

	// tell the sink threads to go
	pthread_mutex_lock(&(p_args->sink_lock));
        pthread_mutex_lock(&(c_args->sink_lock));
	pthread_cond_signal(&(p_args->sink_cond));
	pthread_cond_signal(&(c_args->sink_cond));
	pthread_mutex_unlock(&(p_args->sink_lock));
        pthread_mutex_unlock(&(c_args->sink_lock));

        // wait for the I/O threads to finish
        pthread_join(thread1, NULL);
        pthread_join(thread2, NULL);
	pthread_join(sink_thr1, NULL);
        pthread_join(sink_thr2, NULL);
	
	return 0;

 error_exit3:
	pthread_cancel(thread1);
 error_exit2:
	pthread_cancel(sink_thr2);
 error_exit1:
	pthread_cancel(sink_thr1);
 error_exit:
	return -1;
}

void *psd_xfer_splice_sink(void *arg) {
	psdXferArgs *xfer_args = arg;
        psdConn *curr_sink = xfer_args->sink;
	psdConn *curr_src = xfer_args->src;

	uint64_t sent;
	uint64_t bytes_sent;
        size_t bytes_sent_size;
	struct timeval start_time, end_time;

	int n;
	int ret;
	
	pthread_mutex_lock(&(xfer_args->sink_lock));
	
        psd_info(2, "new sink thread: %s->%s", curr_src->description, curr_sink->description);
	
        pthread_cond_signal(&(xfer_args->sink_cond)); // alert the parent          	

	sent = 0;
	
	gettimeofday(&start_time, NULL);

        while (1) {
		
		pthread_cond_wait(&(xfer_args->sink_cond), &(xfer_args->sink_lock));
		
                if (xfer_args->xfer_amt == (uint64_t) -1) {
                        psd_info(2, "%s->%s: going into perpetual receive mode", curr_src->description, curr_sink->description);
                } else {
                        psd_info(2, "%s->%s: receiving %llu bytes", curr_src->description, curr_sink->description, xfer_args->xfer_amt);
                }
		
                sent = 0;

		while(sent < xfer_args->xfer_amt) {
			
			n = psd_conn_sink_splice(curr_sink, xfer_args->pipes[0], SPLICE_SIZE, 0);
			
			//psdConn_tcpData *tcp_sink = (psdConn_tcpData *) curr_sink->conn_private;
			//n = ssplice(xfer_args->pipes[0], NULL, tcp_sink->sd, NULL, SPLICE_SIZE, 0);
			
			if (n <= 0) {
				
				if (errno == EAGAIN) continue;
				
				if(n < 0) {
					psd_err(0, "%s->%s: sink_splice failed: %s", curr_src->description, curr_sink->description, strerror(errno));
				}
				goto out;
			}                                                                                             
			
			//psd_info(11, "%s->%s: sink_splice returned %d", curr_src->description, curr_sink->description, n);
			
			sent += n;
		}
	}
 out:
	close(xfer_args->pipes[0]);
        close(xfer_args->pipes[1]);

        gettimeofday(&end_time, NULL);

        bytes_sent_size = sizeof(bytes_sent);
	if (psd_conn_get_stat(curr_sink, PSD_STATS_BYTES_WRITTEN, &bytes_sent, &bytes_sent_size) == 0) {
                char tmp[255];
		
                snprintf(tmp, sizeof(tmp), "%s/%s->%s", psd_session_get_id(xfer_args->sess), curr_src->description, curr_sink->description);
		psd_log_measurement_uint(end_time, tmp, "TransferBytes", bytes_sent);
                psd_log_measurement_dbl(end_time, tmp, "TransferRate", ((bytes_sent/difftv(&start_time, &end_time))*8)/1000/1000);
                psd_log_measurement_dbl(end_time, tmp, "TransferTime", difftv(&start_time, &end_time));
        }

	psd_conn_shutdown(curr_sink, PSD_SEND_SIDE);
	pthread_exit(NULL);
}

void *psd_xfer_splice_src(void *arg) {
	psdXferArgs *xfer_args = arg;
	psdConn *curr_src = xfer_args->src;
        psdConn *curr_sink = xfer_args->sink;
	int ret;
	int n;
	uint64_t sent;

	pthread_mutex_lock(&(xfer_args->lock));

	psd_info(2, "new transfer thread: %s->%s", curr_src->description, curr_sink->description);
	
        pthread_cond_signal(&(xfer_args->cond)); // alert the parent                                                                                             
	while (1) {
		
                pthread_cond_wait(&(xfer_args->cond), &(xfer_args->lock));

		if (xfer_args->xfer_amt == (uint64_t) -1) {
                        psd_info(2, "%s->%s: going into perpetual receive mode", curr_src->description, curr_sink->description);
                } else {
                        psd_info(2, "%s->%s: receiving %llu bytes", curr_src->description, curr_sink->description, xfer_args->xfer_amt);
                }
		
		sent = 0;
		
                while(sent < xfer_args->xfer_amt) {
			
			n = psd_conn_src_splice(curr_src, xfer_args->pipes[1], SPLICE_SIZE, 0);
			//psdConn_tcpData *tcp_src = (psdConn_tcpData *) curr_src->conn_private;
			//n = ssplice(tcp_src->sd, NULL, xfer_args->pipes[1], NULL, SPLICE_SIZE, 0);
			
			if (n <= 0) {
				
				if (errno == EAGAIN) continue;

				if(n < 0) {
					psd_err(0, "%s->%s: src_splice failed: %s", curr_src->description, curr_sink->description, strerror(errno));
					
				}
				goto out;
			}
			
			sent += n;

			//psd_info(11, "%s->%s: src_splice returned %d", curr_src->description, curr_sink->description, n);
		}
		
		if (write(xfer_args->parent_fd, "1", 1) < 0) {
                        psd_err(0, "%s->%s: couldn't communicate with parent thread: %s", curr_src->description, curr_sink->description, strerror(errno));
                        goto out;
		}
	}
	

out:
        psd_info(1, "%s->%s: closing sockets and returning.", curr_src->description, curr_sink->description);

	close(xfer_args->pipes[1]);
	close(xfer_args->pipes[0]);

        psd_conn_shutdown(curr_src, PSD_RECV_SIDE);
        pthread_exit(NULL);
}



